package org.flint.data.structure.sgt;

import com.google.common.base.Preconditions;

import java.util.Arrays;

/**
 * @author flint92
 */
public class SegmentTree<E> {

    private E[] tree;

    private E[] elementData;

    private Merge<E> merge;

    @SuppressWarnings("unchecked")
    public SegmentTree(E[] arr, Merge<E> merge) {
        if (arr == null || arr.length == 0) {
            throw new IllegalArgumentException("arr must not be empty");
        }
        int len = arr.length;
        this.elementData = Arrays.copyOf(arr, len);

        this.tree = (E[]) new Object[len << 2];
        this.merge = merge;

        buildSegmentTree(0, 0, len - 1);

    }

    public E get(int index) {
        Preconditions.checkElementIndex(index, elementData.length);
        return elementData[index];
    }

    private int leftChild(int index) {
        // index * 2 + 1
        return index << 1 | 1;
    }

    private int rightChild(int index) {
        // (index + 1) * 2
        return (-(~index)) << 1;
    }


    /**
     * 在treeIndex的位置创建区间[l, r]的线段树
     * @param treeIndex         树的索引
     * @param left              左边界
     * @param right             右边界
     */
    private void buildSegmentTree(int treeIndex, int left, int right) {
        if (left == right) {
            tree[treeIndex] = elementData[left];
            return;
        }

        int leftTreeIndex = leftChild(treeIndex);
        int rightTreeIndex = rightChild(treeIndex);

        int mid = left + ((right - left) >> 1);
        buildSegmentTree(leftTreeIndex, left, mid);
        buildSegmentTree(rightTreeIndex, mid + 1, right);

        tree[treeIndex] = merge.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
    }

    /**
     * 求取区间[queryL, queryR]的值
     * @param queryL        左
     * @param queryR        右
     * @return              区间值
     */
    public E query(int queryL, int queryR) {
        if (queryL < 0 || queryL >= elementData.length
                || queryR < 0 || queryR >= elementData.length
                || queryL > queryR) {
            throw new IllegalArgumentException("Index is illegal.");
        }

        return query(0, 0, elementData.length - 1, queryL, queryR);
    }

    /**
     * 在以treeIndex为根的线段树中[left...right]范围内搜索区间[queryL, queryR]的值
     * @param treeIndex     线段树索引
     * @param left          线段树的左边界
     * @param right         线段树的右边界
     * @param queryL        查询区间左边界
     * @param queryR        查询区间右边界
     * @return              区间值
     */
    private E query(int treeIndex, int left, int right, int queryL, int queryR) {
        if (left == queryL && right == queryR) {
            return tree[treeIndex];
        }
        int mid = left + ((right - left) >> 1);
        int leftTreeIndex = leftChild(treeIndex);
        int rightTreeIndex = rightChild(treeIndex);

        if (queryL >= mid + 1) {
            return query(rightTreeIndex, mid + 1, right, queryL, queryR);
        } else if (queryR <= mid) {
            return query(leftTreeIndex, left, mid, queryL, queryR);
        }

        E leftResult = query(leftTreeIndex, left, mid, queryL, mid);
        E rightResult = query(rightTreeIndex, mid + 1, right, mid + 1, queryR);
        return merge.merge(leftResult, rightResult);
    }


    public void set(int index, E e) {
        Preconditions.checkElementIndex(index, elementData.length);
        elementData[index] = e;

        set(0, 0, elementData.length - 1, index, e);
    }

    /**
     * 在treeIndex为根的线段树中更新index的值为e
     * @param treeIndex     线段树索引
     * @param left          线段树左边界
     * @param right         线段树右边界
     * @param index         更新的索引值
     * @param e             元素值
     */
    private void set(int treeIndex, int left, int right, int index, E e) {
        if (left == right) {
            tree[treeIndex] = e;
            return;
        }

        int mid = left + ((right - left) >> 1);
        int leftTreeIndex = leftChild(treeIndex);
        int rightTreeIndex = rightChild(treeIndex);

        if (index >= mid + 1) {
            set(rightTreeIndex, mid + 1, right, index, e);
        } else {
            set(leftTreeIndex, left, mid, index, e);
        }

        tree[treeIndex] = merge.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
    }
}
